"""
	Copyright 2008 Erik Volz

    This file is part of RDS Simulator.

    RDS Simulator is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    RDS Simulator is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with RDS Simulator.  If not, see <http://www.gnu.org/licenses/>.

erik.volz@gmail.com
http://www.erikvolz.info


initially conceived as a means to calculate the optimal #seeds given 
that sample size should be reached by a desired date

organization
x Parameters
	-off dist -rec intervals -sample size -target due date
	-read/write od/rec times  -#seeds
x RDSSim
- DiGraph -methods for serializing/dot files  
-recruitments[] (rec id, time rec interviewed)
-primary logic of simulation
x RDSSimulationSet
- methods for generating distribution of completion times
- methods for exploring various #seeds etc. 
{possibly in separate package}
-TODO output distribution of path lengths + diameter + mean path length + depth; 
	
long term
classes to explore efficiency of adaptive methods

TODO activation and expiration
"""

from pylab import *
from scipy.stats.distributions import *
import networkx as nx
import time, pdb, csv



class Parameters:
	od = rv_discrete(values=(range(4),[0.40307102,0.21113244,0.17850288,0.20729367]),name='offspring_distribution')
	recIntervals = expon(scale=7.)
	sampleSize = 250
	numSeeds = 8
	targetFinishTime = 75 #days
	recIntervalsScale = 7.
	
	"""def __init__(self, rifn='recruitmentIntervals.dat'):
		self.recIntervals = load(rifn).astype(int)
	#"""
	
	""" assumes an exponential distribution with given exponent """
	def set_recruitment_intervals_exponential(self, r):
		self.recIntervals = expon(scale=1./r)
	#
#

class RDSSim:
	def __init__(self, parameters):
		self.parameters = parameters
		self.parameters.set_recruitment_intervals_exponential(1./self.parameters.recIntervalsScale)
		#self.__dict__.update(parameters.__dict__)
	#
	
	def new_simulation(self):
		#recTree = nx.DiGraph()
		recruitments = list()
		recruitmentsMade = list()
		
		self.tree = nx.DiGraph()
		
		# add seeds to recruitments
		nr = self.__new_recruits(0, 0.)
		recruitments = self.__add_new_recruits(nr, recruitments)
		for i in range(1, self.parameters.numSeeds):
			nr = self.__new_recruits(i, 0.)
			recruitments = self.__add_new_recruits(nr, recruitments)
		#
		self.tree.add_nodes_from(range(self.parameters.numSeeds) , depth=0)
		
		nn = self.parameters.numSeeds
		n = self.parameters.sampleSize
		while nn < n:
			# check that sample has not died out
			if len(recruitments)==0:
				break
			#pick min value from queue
			nextRecruitment = recruitments.pop()
			recruitmentsMade.append(nextRecruitment)
			# make new recruits
			nr = self.__new_recruits(nn, nextRecruitment[1])
			# add recruits of new recruiter to recruitments
			recruitments = self.__add_new_recruits(nr, recruitments)
			
			# update digraph
			self.tree.add_edge(nextRecruitment[0], nn)
			self.tree.node[nn]['depth'] = self.tree.node[nextRecruitment[0]]['depth'] + 1
			
			
			nn+=1
		#
		
		#self.recTree = recTree
		if nn==n:
			self.sampleSizeReached = True
			self.finishTime = nextRecruitment[1]
			self.sampleSizeReachedOnTarget= self.finishTime<self.parameters.targetFinishTime
		else:
			self.sampleSizeReached = False
			self.sampleSizeReachedOnTarget = False
			try:
				self.finishTime = nextRecruitment[1]
			except: #seeds made no recruitments!
				self.finishTime = 0.
	#
	
	def __new_recruits(self, recruiterID, baseTime):
		nr = list()
		for i in range(int(self.parameters.od.rvs())):
			#ri = self.parameters.recIntervals[randint.rvs(0,len(self.parameters.recIntervals)-1)]
			ri = self.parameters.recIntervals.rvs()
			t = baseTime + ri
			nr.append((recruiterID, t))
		#
		return nr
	#
	
	def __add_new_recruits(self, newrecs, oldrecs):
		# insert in proper order
		# obey sample size limit
		#TODO this code could be much more efficient
		for rec in newrecs:
			if len(oldrecs)==1:
				if rec[1] > oldrecs[0][1]:#belongs at the beginning
					oldrecs.insert(0, rec)
				else:
					oldrecs.append(rec)
				#
			#
			elif len(oldrecs)>1:
				if rec[1] > oldrecs[0][1]:#belongs at the beginning
					oldrecs.insert(0, rec)
				#
				else:
					for i in range(1, len(oldrecs)): # belongs in the middle
						if rec[1] > oldrecs[i][1] and rec[1] < oldrecs[i-1][1]:
							oldrecs.insert(i, rec)
							break;
						#
					#
					if rec[1]<oldrecs[-1][1]: #belongs at the end
						oldrecs.append(rec)
					#
				#
			elif len(oldrecs)==0:
				oldrecs.append(rec) # if empty, initialize with rec
			#
		#
		if len(oldrecs) > self.parameters.sampleSize:
			oldrecs = oldrecs[-self.parameters.sampleSize:]
		#
		return oldrecs # which are now new
	#
	
	def max_depth(self):
		if not self.__dict__.has_key('tree'):
			return None
		return max( nx.get_node_attributes(self.tree, 'depth').values() )
	def average_depth(self):
		if not self.__dict__.has_key('tree'):
			return None
		return mean( nx.get_node_attributes(self.tree, 'depth').values() )
	def mean_path_length(self):
		if not self.__dict__.has_key('tree'):
			return None
		#~ introduce source for seeds; then samples with many seeds are appropriately penalized
		utree = self.tree.to_undirected()
		utree.add_edges_from(  nx.complete_graph(self.parameters.numSeeds).edges() )
		return nx.average_shortest_path_length(utree)
	#~ def max_path_length(self):
		#~ if not self.__dict__.has_key('tree'):
			#~ return None
		#~ 
	#
	
	def outputGraphVizTree(self, outfn):
		import pygraphviz as pgv
		rectree = nx.to_agraph(self.tree)
		#~ rectree.graph_attr['label']='Name of graph'
		rectree.node_attr['shape']='circle'
		rectree.node_attr['fillcolor']='black'
		rectree.edge_attr['color']='black'
		for u in rectree.nodes():
			u.attr['label'] = ''
			u.attr['style'] = 'filled'
			u.attr['fillcolor'] = 'black'
			u.attr['color'] = 'black'
		#
		rectree.layout(prog='neato')
		rectree.draw(outfn)
		# for visualizing trees:
		"""node = recTree.get_node(nn)
		node.attr['width']=repr(nextRecruitment[1]/100)
		node.attr['height']=repr(nextRecruitment[1]/100)
		node.attr['fixedsize'] = 'true'"""
		#edge = recTree.get_edge(nextRecruitment[0],nn)
		#edge.attr['len']=repr(nextRecruitment[1])
		#
	#
#

"""
author: Erik Volz
date July 25, 2008
"""

class SimulationSet:
	numSimulations = 200
	odvalues = [0.40307102,0.21113244,0.17850288,.20729367]
	od =rv_discrete(values=(range(4),odvalues),name='offspring_distribution')

	
	sampleSize = [250]
	numSeeds = [8]
	recIntervalsScale = [7.]
	
	variableParameters = ['numSeeds', 'recIntervalsScale', 'sampleSize']
	niceVariableNames = ['Number of Seeds', 'Recruitment Interval Scale', 'Sample Size']
	# static variables includ offspring distribution, targetFinishTime
	
	def __init__(self, **args):
		"""# to prevent static/dynamic errors, static 
		definitions should be repeated here """
		self.sampleSize = [250]
		self.numSeeds = [8]
		self.recIntervalsScale = [7.]
		self.targetFinishTime = 75
		
		for k in args:
			self.__dict__[k] = args[k]
		#
	#
	
	""" parameter combinations serve as keys for the dictionaries
	which hold output; each refers to the indices of the corresponding
	parameter values  
	
	given that there are only three variable parameters, this is a little
	bit of over-kill, but it insures that the code can be scaled up for more
	variable parameters in
	the future if need be"""
	def __combo_from_integer(self, i, pp):
		combo = list()
		ii = i
		for k in range(1,len(pp)+1):
			j = ii / pp[-k]
			combo.append(j)
			ii = ii-j*pp[-k]
		#
		return tuple(combo[::-1])
		#
	def __parameter_combinations(self):
		#periodLengths  = [len(self.sampleSize), len(self.numSeeds), len(self.recIntervalsScale)]
		try:
			periodLengths = [len(self.__dict__[x]) for x in self.variableParameters ]
		except: # stupid dynamic/static dictionary problems
			periodLengths = list()
			for i in range(len(self.variableParameters)):
				vn = self.variableParameters[i]
				try:
					periodLengths.append(len(self.__dict__[vn]))
				except:
					periodLengths.append(len(SimulationSet.__dict__[vn]))
				#
			#
		#
		partialProd = [1]
		for i in range(1, len(periodLengths)):
			partialProd.append(partialProd[i-1]*periodLengths[i-1])
		#
		totalNumCombos = prod(periodLengths)
		combos = list()
		for i in range(totalNumCombos):
			combos.append( self.__combo_from_integer(i, partialProd))
		#
		return combos
	####
	
	
	def collect_set(self):
		parameterCombinations = self.__parameter_combinations()
		#pdb.set_trace()
		ft = dict() # finish times
		ssr = dict() # sample size reached boolean
		ssrot = dict() # sample size reached on target
		maxdepth = self.maxdepth = dict()
		meandepth = self.meandepth =dict()
		meanpathlength = self.meanpathlength = dict()
		#sims = list()
		for combo in parameterCombinations:
			p = Parameters()
			for i in range(len(combo)):
				vn = self.variableParameters[i]
				try:
					p.__dict__[vn] = self.__dict__[vn][combo[i]]
				except:
					try:
						p.__dict__[vn] = SimulationSet.__dict__[vn][combo[i]]
					except:
						pdb.set_trace()
			#
			p.od = self.od #
			p.targetFinishTime = self.targetFinishTime
			ft[combo] = list()
			ssr[combo] = list()
			ssrot[combo] = list()
			maxdepth[combo] = list()
			meandepth[combo] = list()
			meanpathlength[combo] = list()
			s = RDSSim(p)
			for i in range(self.numSimulations):
				s.new_simulation()
				ft[combo].append( s.finishTime)
				ssr[combo].append( s.sampleSizeReached)
				ssrot[combo].append(s.sampleSizeReachedOnTarget)
				maxdepth[combo].append(s.max_depth())
				meandepth[combo].append(s.average_depth())
				meanpathlength[combo].append(s.mean_path_length())
			#sims.append(s)
		#
		self.finishTimes = ft
		self.sampleSizeReached = ssr
		#self.sims = sims
		self.sampleSizeReachedOnTarget = ssrot
		
		return True
	#
	
	
	# DONE should not print key, but variable combination
	# DONE should output % reached by target date
	def print_output(self):
		s = ""
		for k in self.finishTimes.keys():
			s+= 'key:'+repr(k)+'\n'
			for i in range(len(k)):
				s+=self.niceVariableNames[i]+':'+repr(self.__dict__[self.variableParameters[i]][k[i]])+'\n'
			#
			s+= 'Mean finish time:'+ repr(mean(self.finishTimes[k]))+'\n'
			#print 'finish times:', self.finishTimes[k]
			#~ s+= 'Mean time to recruitment:' + self. + '\n'
			fssr = mean(self.sampleSizeReached[k])
			fssrot = mean(self.sampleSizeReachedOnTarget[k])
			maxdepth = mean(self.maxdepth[k])
			meandepth = mean(self.meandepth[k])
			mpl = mean(self.meanpathlength[k])
			s+= 'Fraction sample size reached: %(#).3f' % {'#': fssr} +'\n'
			s+= 'Fraction sample reached by target date: %(#).3f' % { '#': fssrot } +'\n'
			s+= 'Maximum recruitment tree depth: %1.2f \n' % maxdepth
			s+= 'Mean recruitment tree depth: %1.2f \n' % meandepth
			s+= 'Mean path length within recruitment trees: %1.2f \n' % mpl
			s+= '##################################################\n'
			#figure(k)
			#hist(self.finishTimes[k]); show()
		# 
		print s
		return s
	#
	
	def save_csv(self, fn):
		labels = ['NumberSeeds', 'RecruitmentScale','SampleSize','FinishTime','CompleteByTargetFinishTime','Complete']
		data = list()
		for k in self.finishTimes.keys():
			ns = self.numSeeds[k[0]]
			ris = self.recIntervalsScale[k[1]]
			ss = self.sampleSize[k[2]]
			data.extend([[ns,ris,ss,self.finishTimes[k][i],self.sampleSizeReachedOnTarget[k][i],self.sampleSizeReached[k][i]] for i in range(len(self.finishTimes[k])) ])
		#
		
		f = open(fn,'w')
		w = csv.writer(f)
		w.writerow(labels)
		w.writerows(data)
		f.close()
	#
	""" DONE need an output method that takes two variable parameters
	and one fixed parameter value
	and generates  3D figure
	with vertical axis giving either a) % complete by finish time
	b) mean finish time c)% complete at all
	"""
# TESTING
if __name__ == '__main__':
	"""
	p = Parameters() # default parm's
	p.set_recruitment_intervals_exponential(1./7)
	s = RDSSim(p)
	s.new_simulation()
	s.outputGraphVizTree('seeds8')"""
	
	"""print time.ctime()
	ss = SimulationSet()
	ss.collect_set_by_seed_values()
	ss.print_output()
	print time.ctime() """
	
	"""print 'For STAHR'
	print time.ctime()
	ss = SimulationSet()
	ss.numSimulations = 1000
	ss.numSeeds = [8, 15]
	ss.recIntervalsScale = [7., 14.]
	ss.collect_set()
	ss.print_output()
	print time.ctime()"""
	
	print 'DEBUG'
	print time.ctime()
	ss = SimulationSet()
	ss.numSimulations = 5
	ss.numSeeds = [8, 10, 30]
	ss.recIntervalsScale = [7., 100.]
	ss.collect_set()
	ss.print_output()
	ss.save_csv('test.csv')
	print time.ctime()
#

